{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Licensed under the Apache License, Version 2.0"
      ],
      "metadata": {
        "id": "HBt-ZDa8v-BN"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Example code for reprojecting GOES-16 images\n",
        "This notebook demonstrates how to reproduce the reprojected images in the OpenContrails dataset."
      ],
      "metadata": {
        "id": "dvdYHr9wvI6z"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install pyresample\n",
        "!pip install gcsfs\n",
        "!pip install xarray"
      ],
      "metadata": {
        "id": "fpUkh7s3lSaY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qZF3Ypk7dTwq"
      },
      "outputs": [],
      "source": [
        "import datetime\n",
        "import pprint\n",
        "import sys\n",
        "\n",
        "import gcsfs\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "from osgeo import osr\n",
        "import pyresample\n",
        "import pyresample.bilinear\n",
        "import tensorflow as tf\n",
        "import xarray\n",
        "\n",
        "if 'google.colab' in sys.modules:\n",
        "  from google.colab import auth\n",
        "  auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Load a single record from the TFRecords and display the 12um band.\n",
        "\n",
        "def parse_example(serialized_example: bytes) -\u003e dict[str, tf.Tensor]:\n",
        "  features = tf.io.parse_single_example(serialized_example, {\n",
        "      'cloud_top':  tf.io.FixedLenFeature([], tf.string),\n",
        "      'data_10um':  tf.io.FixedLenFeature([], tf.string),\n",
        "      'data_11um':  tf.io.FixedLenFeature([], tf.string),\n",
        "      'data_12um':  tf.io.FixedLenFeature([], tf.string),\n",
        "      'human_pixel_masks': tf.io.FixedLenFeature([], tf.string),\n",
        "      'n_times_before': tf.io.FixedLenFeature([], tf.int64),\n",
        "      'n_times_after': tf.io.FixedLenFeature([], tf.int64),\n",
        "      # Projection params\n",
        "      'projection_wkt': tf.io.FixedLenFeature([], tf.string),\n",
        "      'col_min': tf.io.FixedLenFeature([], tf.float32),\n",
        "      'row_min': tf.io.FixedLenFeature([], tf.float32),\n",
        "      'col_size': tf.io.FixedLenFeature([], tf.float32),\n",
        "      'row_size': tf.io.FixedLenFeature([], tf.float32),\n",
        "      # Timestamp\n",
        "      'timestamp': tf.io.FixedLenFeature([], tf.int64),  # approximate timestamp\n",
        "      'satellite_scan_starts': tf.io.VarLenFeature(tf.int64),  # timestamp from original file\n",
        "  })\n",
        "  for key in ['cloud_top', 'data_10um', 'data_11um', 'data_12um']:\n",
        "    features[key] = tf.io.parse_tensor(features[key], tf.double)\n",
        "  features['human_pixel_masks'] = tf.io.parse_tensor(features['human_pixel_masks'], tf.int32)\n",
        "  return features\n",
        "\n",
        "dataset = tf.data.TFRecordDataset(tf.io.gfile.glob('gs://goes_contrails_dataset/20230419/tfrecords/train.tfrecords-*'))\n",
        "dataset = dataset.map(parse_example)\n",
        "features = dataset.take(1).get_single_element()\n",
        "\n",
        "n_times_before = features['n_times_before']\n",
        "plt.figure(figsize=(12, 6))\n",
        "plt.imshow(features['data_12um'][:, :, n_times_before])\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "y9RJIjajUxcC",
        "collapsed": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here we load the original NetCDF files that is publicly available from Google Cloud Storage. We convert the raw radiance to brightness temperature, and then compute the `AreaDefinition` of the original GOES-16 full-disk image from the parameters in the NetCDF files. \n",
        "\n",
        "See more in the [example  notebook](https://github.com/google-research/google-research/blob/master/contrails/demos/load_goes_data.ipynb) for loading and visualizing GOES images."
      ],
      "metadata": {
        "id": "szQLfMNRqc2V"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "satellite_scan_starts = tf.sparse.to_dense(features['satellite_scan_starts']).numpy()\n",
        "timestamp = satellite_scan_starts[features['n_times_before']]\n",
        "dt = datetime.datetime.fromtimestamp(timestamp, tz=datetime.timezone.utc)\n",
        "\n",
        "fs = gcsfs.GCSFileSystem(project='gcp-public-data-goes-16')\n",
        "\n",
        "# data_12um corresponds to channel 15 (i.e. M6C15)\n",
        "paths = fs.glob(f'gcp-public-data-goes-16/ABI-L1b-RadF/{dt.year}/{dt.timetuple().tm_yday:03d}/{dt.hour:02d}/OR_ABI-L1b-RadF-M6C15_G16_s{dt.year}{dt.timetuple().tm_yday:03d}{dt.hour:02d}{dt.minute:02d}*')\n",
        "assert len(paths) == 1, 'There should be exactly one NetCDF file for a band at the timestamp.'\n",
        "\n",
        "with fs.open(paths[0], 'rb') as f:\n",
        "  dataset = xarray.open_dataset(f)\n",
        "  dataset.load()\n",
        "\n",
        "# Convert the raw radiance to brightness temperature.\n",
        "radiance = dataset.Rad.data\n",
        "brightness_temperature = (dataset.planck_fk2.data / np.log((dataset.planck_fk1.data / radiance) + 1) - dataset.planck_bc1.data) / dataset.planck_bc2.data"
      ],
      "metadata": {
        "id": "Qr5ALvB_JQCZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "h0 = dataset.goes_imager_projection.perspective_point_height\n",
        "goes_area_def = pyresample.geometry.AreaDefinition(\n",
        "  area_id='all_goes_16',  # Used only for pyresample logging\n",
        "  proj_id='deprecated',  # Deprecated but required by pyresample\n",
        "  description='all_goes_16',  # Used only for pyresample logging\n",
        "  projection={  # proj4 dict\n",
        "      'proj': 'geos',  # Stands for 'geostationary'\n",
        "      'units': 'm',\n",
        "      'h': str(h0),\n",
        "      'lon_0': str(\n",
        "          dataset.goes_imager_projection.longitude_of_projection_origin\n",
        "      ),\n",
        "      'a': str(dataset.goes_imager_projection.semi_major_axis),\n",
        "      'b': str(dataset.goes_imager_projection.semi_minor_axis),\n",
        "      'sweep': dataset.goes_imager_projection.sweep_angle_axis,\n",
        "  },\n",
        "  width=dataset['x'].shape[0],\n",
        "  height=dataset['y'].shape[0],\n",
        "  area_extent=[\n",
        "      dataset['x_image_bounds'].data[0] * h0,\n",
        "      dataset['y_image_bounds'].data[1] * h0,\n",
        "      dataset['x_image_bounds'].data[1] * h0,\n",
        "      dataset['y_image_bounds'].data[0] * h0,\n",
        "  ],\n",
        ")\n"
      ],
      "metadata": {
        "id": "HqOKil2Vqk1_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here we construct the `AreaDefinition` for the GOES scene from the projection parameters provided in the TFRecords."
      ],
      "metadata": {
        "id": "QGhN-eMuqPi5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "rows, cols = features['data_12um'].shape[:2]\n",
        "area_extent = [\n",
        "    features['col_min'],\n",
        "    features['row_min'] + features['row_size'] * rows,\n",
        "    features['col_min'] + features['col_size'] * cols,\n",
        "    features['row_min']\n",
        "]\n",
        "target_area_def = pyresample.AreaDefinition(\n",
        "    area_id='n/a',\n",
        "    description='n/a',\n",
        "    proj_id='n/a',\n",
        "    projection=osr.SpatialReference(wkt=features['projection_wkt'].numpy().decode()).ExportToProj4(),\n",
        "    width=cols,\n",
        "    height=rows,\n",
        "    area_extent=area_extent,\n",
        ")"
      ],
      "metadata": {
        "id": "Kv9SHMHa_Q3h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Once we have the `AreaDefinition` of the original GOES full disk image and the target scene, we can use bilinear resampling to obtain the image that corresponds to the local scene."
      ],
      "metadata": {
        "id": "xBUkoeFmqqxI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "t_params, s_params, input_idxs, idx_ref = pyresample.bilinear.get_bil_info(\n",
        "      goes_area_def, target_area_def\n",
        ")\n",
        "resampled = pyresample.bilinear.get_sample_from_bil_info(\n",
        "    brightness_temperature.flatten(),\n",
        "    t_params,\n",
        "    s_params,\n",
        "    input_idxs,\n",
        "    idx_ref,\n",
        "    output_shape=target_area_def.shape,\n",
        ")"
      ],
      "metadata": {
        "id": "GuPjWy16LPzx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The resampled image closely reproduces the one provided with the dataset."
      ],
      "metadata": {
        "id": "E556qjVgq4_1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "plt.figure(figsize=(12, 6))\n",
        "plt.subplot(1, 2, 1)\n",
        "plt.imshow(resampled)\n",
        "plt.subplot(1, 2, 2)\n",
        "plt.imshow(features['data_12um'][..., features['n_times_before']])"
      ],
      "metadata": {
        "id": "vmbnB7KwLr9f"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
